import os.path

import paddle
import paddlehub as hub

from src.config.config import DATA_ROOT, LABEL_PATH, TRAIN_PATH
from src.data.custom_cv_dataset import CustomCVDataset

# 加载预训练模型
module = hub.Module(name="resnet50_vd_animals")
# # 修改输出层适配新任务（如新类别数）
# with hub.finetune.strategy(parameters=module.parameters()):
#     module.add_config(
#         num_classes=10,  # 新任务的类别数
#         label_map={"cat": 0}  # 自定义标签映射
#     )

path = os.path.join(DATA_ROOT, '一级动物')
dataset = CustomCVDataset(base_path=path, train_list_file=TRAIN_PATH, label_list_file=LABEL_PATH)

# 配置训练参数
optimizer = paddle.optimizer.Adam(learning_rate=1e-4,  parameters=module.parameters())
# 6. 创建Trainer
trainer = hub.Trainer(
    module=module,
    optimizer=optimizer,
    use_gpu=True,
    dataset=dataset,
    batch_size=32,
    num_epochs=10,
    eval_interval=200  # 每200步验证一次
)

# 7. 启动训练
trainer.train()

hub.save(module,  "finetuned_resnet50_animals")
